from osgeo import gdal
import pandas as pd
from mpl_toolkits.basemap import Basemap
import numpy as np
import matplotlib.pyplot as plt
import os
from scipy.interpolate import griddata


def plot_paddy_field(tif_path, sample_mark, output_path, mark):
    mark_data = pd.read_csv(sample_mark)
    selct_data = mark_data.loc[mark_data['Class_Level2'] == mark]
    in_ds = gdal.Open(tif_path)
    adfGeoTransform = in_ds.GetGeoTransform()
    nXSize = in_ds.RasterXSize
    nYSize = in_ds.RasterYSize
    im_data = in_ds.ReadAsArray(0, 0, nXSize, nYSize)
    im_data_ravel = im_data.ravel()
    lon_px = []
    lat_py = []
    for i in range(nYSize):
        for j in range(nXSize):
            px = adfGeoTransform[0] + j * adfGeoTransform[1]
            py = adfGeoTransform[3] + i * adfGeoTransform[5]
            lon_px.append(px)
            lat_py.append(py)
    data_griddata = griddata((lon_px, lat_py), im_data_ravel, (selct_data["XLon"], selct_data["YLat"]), method='linear')
    plt.figure(figsize=[12, 6])
    plt.title('Paddy field distribution')
    map = Basemap(projection='cyl', resolution='l')
    map.drawcoastlines()
    map.drawcountries()
    map.drawmapboundary()
    lon = np.array(selct_data["XLon"])
    lat = np.array(selct_data["YLat"])
    map.scatter(lon, lat, c=data_griddata, s=10)
    # plt.show()
    plt.savefig(output_path + "_paddy_field.jpg")


if __name__ == '__main__':
    mark = 11
    tif_path = 'F:/16satellite/data/merge/2007/001/A2007_001_global_normalization.tif'
    sample_mark = 'F:/16satellite/data/sample_mark/sample_mark.csv'
    (filename, extension) = os.path.splitext(tif_path)
    plot_paddy_field(tif_path, sample_mark, filename, mark)

